import cv2
import numpy as np


def combine_images(images):  # 滤波后图像与频域图组合在一起
    shapes = np.array([mat.shape for mat in images])
    rows = np.max(shapes[:, 0])
    copy_imgs = [cv2.copyMakeBorder(img, 0, rows - img.shape[0], 0, 0,
                                    cv2.BORDER_CONSTANT, (0, 0, 0)) for img in images]
    return np.hstack(copy_imgs)


def fft(img):  # 傅里叶变换
    rows, cols = img.shape[:2]
    nrows = cv2.getOptimalDFTSize(rows)  # 得到傅里叶最优尺寸大小
    ncols = cv2.getOptimalDFTSize(cols)
    nimg = np.zeros((nrows, ncols))
    nimg[:rows, :cols] = img
    fft_mat = cv2.dft(np.float32(nimg), flags=cv2.DFT_COMPLEX_OUTPUT)
    return np.fft.fftshift(fft_mat)


def fft_image(fft_mat):
    log_mat = cv2.log(1 + cv2.magnitude(fft_mat[:, :, 0], fft_mat[:, :, 1]))
    cv2.normalize(log_mat, log_mat, 0, 255, cv2.NORM_MINMAX)
    return np.uint8(np.around(log_mat))


def ifft(fft_mat):  # 逆傅里叶变换
    f_ishift_mat = np.fft.ifftshift(fft_mat)
    img_back = cv2.idft(f_ishift_mat)
    img_back = cv2.magnitude(*cv2.split(img_back))
    cv2.normalize(img_back, img_back, 0, 255, cv2.NORM_MINMAX)
    return np.uint8(np.around(img_back))


def fft_distances(m, n):
    u = np.array([i if i <= m / 2 else m - i for i in range(m)], dtype=np.float32)
    v = np.array([i if i <= m / 2 else m - i for i in range(m)], dtype=np.float32)
    v.shape = n, 1
    ret = np.sqrt(u * u + v * v)
    return np.fft.fftshift(ret)


def BWfilter(rows, cols, d0, n):  # 巴特沃斯低通滤波
    duv = fft_distances(*fft_mat.shape[:2])
    filter_mat = 1 / (1 + np.power(duv / d0, 2 * n))
    filter_mat = cv2.merge((filter_mat, filter_mat))
    return filter_mat


def do_filter(_=None):
    d0 = cv2.getTrackbarPos('D0', filter_win)
    n = cv2.getTrackbarPos('n', filter_win)
    filter_mat = BWfilter(fft_mat.shape[0], fft_mat.shape[1], d0, n)
    filtered_mat = filter_mat * fft_mat
    img_back = ifft(filtered_mat)
    cv2.imshow(image_win, combine_images([img_back, fft_image(filter_mat)]))


if __name__ == '__main__':
    img = cv2.imread(r"C:\Users\Public\opencv\Figure\lena.jpg", 0)
    rows, cols = img.shape[:2]
    filter_win = 'Filter Parameters'
    image_win = 'Butterworth Low Pass Filtered Image'
    cv2.namedWindow(filter_win)
    cv2.namedWindow(image_win)
    fft_mat = fft(img)
    cv2.createTrackbar('D0', filter_win, 20, min(rows, cols) // 4, do_filter)
    cv2.createTrackbar('n', filter_win, 1, 5, do_filter)
    do_filter()
    cv2.resizeWindow(filter_win, 512, 20)
    cv2.waitKey(0)
    cv2.destroyAllWindows()